import numpy as np
from cv2 import cv2
import re
import matplotlib.pyplot as plt
from utils2 import *
import random
from sklearn.cluster import KMeans
from sklearn import decomposition
import matplotlib
import matplotlib.image as gImage
from sklearn.manifold import TSNE
from matplotlib.ticker import FuncFormatter
import scipy.stats
import time
import random
from sklearn.metrics import confusion_matrix
import copy
import pickle

training_round = 2
query_round = 1

index_map_file = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/IndexMap.pickle'
training_ds = np.load('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round%d/round%d_%dds.npy'%(training_round, training_round, training_round))
query_ds = np.load('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round%d/round%d_%dds.npy'%(training_round, training_round, query_round))
f = open(index_map_file,'rb')
mapping = pickle.load(f)
mapping = mapping[training_round-1]

# sample_idx = 31621
# plt.figure()
# plt.bar(range(100), query_ds[sample_idx], width=1, align='edge')
# plt.title('%d'%sample_idx)
# plt.ylim(0,0.08)
# plt.show()

js_d = []
training_frame_cnt = np.shape(training_ds)[0]
for i in range(training_frame_cnt):
    query_idx = mapping[i][query_round-1]
    if query_idx == -1:
        query_idx = i
    ds1 = training_ds[i]
    ds2 = query_ds[query_idx]
    js_d.append(JS_D(ds1,ds2))

js_d = np.array(js_d)
plt.figure()
plt.scatter(range(training_frame_cnt), js_d, s = 1)
plt.title('JS-D of each frame between training round %d and query round %d'%(training_round, query_round))
plt.show()
